#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include "control.cpp"

namespace py = pybind11;

int add(int i, int j)
{
    return i + j;
}

PYBIND11_MODULE(cworld, m)
{
    m.doc() = "pybind11 example plugin"; // optional module docstring

    m.def("add", &add, "A function which adds two numbers");
    m.def("reset", &reset, "重置环境，并输出初始状态");
    m.def("step", &step, "重置环境，并输出初始状态");
    pybind11::class_<SmartBall>(m, "SmartBall")
        .def("get_obs", &SmartBall::get_obs)
        .def("get_reward", &SmartBall::get_reward)
        .def("get_done", &SmartBall::get_done)
        .def("get_id", &SmartBall::get_id);
}